from kubernetes import client, config
from datetime import datetime
import pytz

class KubernetesHelper:
    def __init__(self, config_file=None):
        # 加载 Kubernetes 配置
        if config_file:
            config.load_kube_config(config_file=config_file)
        else:
            config.load_kube_config()

        # 创建 Kubernetes API 客户端
        self.v1 = client.CoreV1Api()

    def restart_pod(self, namespace, name, container_name):
        # 重启 Pod
        body = {'spec': {'containers': [{'name': container_name, 'restartCount': 0}]}}
        self.v1.patch_namespaced_pod(name, namespace, body)

    def delete_pod(self, namespace, name):
        # 删除 Pod
        self.v1.delete_namespaced_pod(name, namespace)

    def list_pods(self, namespaces=None):
        if namespaces is None:
            namespaces = ['default']  # 默认为default命名空间

        for namespace in namespaces:
            pod_list = self.v1.list_namespaced_pod(namespace, watch=False)

            # 输出命名空间名称
            print(f'NAMESPACE: {namespace}')
            print(f'{"NAME":30} {"READY":10} {"STATUS":15} {"RESTARTS":>12} {"AGE":>5}')

            # 遍历 Pod 列表
            for pod in pod_list.items:
                # 获取所有容器状态信息
                container_statuses = pod.status.container_statuses
                if container_statuses is not None:
                    ready_count = 0
                    for status in container_statuses:
                        if status.ready:
                            ready_count += 1

                        if status.state.waiting is not None or status.state.terminated is not None:
                            pod_status = 'Error'
                            break
                    else:
                        pod_status = 'Running'

                    restart_count = container_statuses[0].restart_count if container_statuses else 0
                else:
                    ready_count = 0
                    pod_status = 'Pending'
                    restart_count = 0

                created_at = pod.metadata.creation_timestamp.timestamp()
                current_time = datetime.now(pytz.utc)
                age = round((current_time.timestamp() - created_at) / 3600, 1)

                # 打印输出每个 Pod 的名称、状态、重启次数和年龄
                print(f'{pod.metadata.name:30} {ready_count}/{len(container_statuses):<4} {pod_status:15} {restart_count:>12} {age:>3.1f}h')

            # 输出分隔符
            print('-' * 80)
    
    def list_all_services(self, namespaces):
        # 遍历每个命名空间下的 Service 列表，并输出分隔行和表头
        for namespace in namespaces:
            print('-' * 80)
            print(f'Namespace: {namespace}\n')
            print(f'{"NAME":30} {"TYPE":10} {"CLUSTER-IP":20} {"EXTERNAL-IP(S)":20} {"PORT(S)":20} {"AGE":15}')

            service_list = self.v1.list_namespaced_service(namespace, watch=False)

            # 遍历 Service 列表
            for service in service_list.items:
                # 获取 Cluster IP 和 Node Port
                cluster_ip = service.spec.cluster_ip or ''
                node_port = ''
                if service.spec.type == 'NodePort':
                    ports = [f'{port.node_port}/{port.protocol}' for port in service.spec.ports]
                    node_port = ', '.join(ports)

                # 获取 External IPs
                external_ips = []
                if service.spec.external_name:
                    external_ips.append(service.spec.external_name)
                if service.spec.load_balancer_ip:
                    external_ips.append(service.spec.load_balancer_ip)
                if service.spec.load_balancer_source_ranges:
                    external_ips.extend(service.spec.load_balancer_source_ranges)
                external_ips = ', '.join(external_ips)

                # 获取年龄
                age = service.metadata.creation_timestamp

                # 打印输出每个 Service 的名称、类型、Cluster IP、External IPs、Ports 和年龄
                print(
                    f'{service.metadata.name:30} {service.spec.type:10} {cluster_ip:20} {external_ips:20} {node_port:20} {age:15}')

    
    def create_pod(self, namespace, pod_name, container_name, image, container_port, host_port):
        # 创建 Pod
        pod_manifest = {
            'apiVersion': 'v1',
            'kind': 'Pod',
            'metadata': {
                'name': pod_name
            },
            'spec': {
                'containers': [{
                    'name': container_name,
                    'image': image,
                    'ports': [{
                        'containerPort': container_port,
                        'hostPort': host_port
                    }]
                }]
            }
        }
        resp = self.v1.create_namespaced_pod(body=pod_manifest, namespace=namespace)
        return resp

    def create_service(self, namespace, service_name, selector_app, port, target_port, node_port):
        # 创建 Service
        service_manifest = {
            'apiVersion': 'v1',
            'kind': 'Service',
            'metadata': {
                'name': service_name
            },
            'spec': {
                'selector': {
                    'app': selector_app
                },
                'ports': [{
                    'protocol': 'TCP',
                    'port': port,
                    'targetPort': target_port,
                    'nodePort': node_port
                }],
                'type': 'NodePort'
            }
        }
        api_instance = client.CoreV1Api()
        resp = api_instance.create_namespaced_service(body=service_manifest, namespace=namespace)
        return resp

    def create_pod_and_service(self, namespace, pod_name, container_name, image, container_port, host_port,
                               service_name, selector_app, port, target_port, node_port):
        # 创建 Pod
        pod_resp = self.create_pod(namespace, pod_name, container_name, image, container_port, host_port)
        print("Pod created. status='%s'" % pod_resp.status.phase)

        # 创建 Service
        service_resp = self.create_service(namespace, service_name, selector_app, port, target_port, node_port)
        print("Service created. status='%s'" % service_resp)

    
# 示例用法
if __name__ == "__main__":
    helper = KubernetesHelper(config_file='/tmp/config')
    
    namespace = 'my-test'
    pod_name = 'mypod'
    container_name = 'mycontainer'
    image = 'nginx:latest'
    container_port = 80
    host_port = 80
    service_name = 'myservice'
    selector_app = 'myapp'
    port = 80
    target_port = 80
    node_port = 30001

    # 创建 Pod 和 Service
    helper.create_pod_and_service(namespace, pod_name, container_name, image, container_port, host_port,
                                   service_name, selector_app, port, target_port, node_port)
    # 重启容器
    helper = KubernetesHelper()
    namespace = 'your-namespace'
    name = 'your-pod-name'
    container_name = 'your-container-name'
    helper.restart_pod(namespace, name, container_name)
    
    # 删除容器
    helper = KubernetesHelper(config_file='/root/.kube/config')
    namespace = 'default'
    pod_list = helper.get_pods(namespace)
    for pod in pod_list:
        print(f'{pod.metadata.namespace}/{pod.metadata.name}')

    helper = KubernetesHelper(config_file='config')
    # 查看默认命名空间下的 Pod
    helper.list_pods()
    # 查看多个命名空间下的 Pod
    helper.list_pods(namespaces=['my-test', 'opx'])
        
    # 查看所有service
    helper = KubernetesHelper(config_file='/tmp/config')
    namespaces = ['my-test', 'opx', 'default']  # 命名空间名称列表
    helper.list_all_services(namespaces)
